package lex.dm.id3;

import java.util.*;

import javax.swing.JTree;

import org.lex.swing.Swings;
import org.lex.utils.Pair;

import lex.dm.utils.Utils;

public class ID3
{
	private static final String ERROR = "(error)";
	/**
	 * �ɼ�������
	 */
	private boolean[] visable;
	private Object[][] data;
	private int classAttrIndex;
	// -------------------
	private IDataSet dataSet = null;

	public static void main( String[] args )
	{
		DataSetUtils.dump( ArrayDataSet.getSampleDataSet(), System.out );
		ID3 id3 = new ID3( ArrayDataSet.getSampleDataSet() );
		JTree tree = new JTree( id3.createTree() );
		Swings.expandAll( tree );
		Swings.show( tree );
	}

	public ID3( IDataSet ds )
	{
		this.dataSet = ds;
	}

	public ID3Node createTree()
	{
		// ��ݳ�ʼ��
		int index = dataSet.getClassAttributeIndex();
		this.data = dataSet.getData();
		this.classAttrIndex = index;
		this.visable = new boolean[this.data[0].length];
		Arrays.fill( this.visable, false );
		this.visable[index] = true;
		// ������������ڵ�
		return this.createDTree();
	}

	private ID3Node createDTree()
	{
		// Object[] maxgain = getMaxGain( array );
		int maxGainIndex = this.getMaxGain( this.data ).second();
		ID3Node root = new ID3Node();
		root.parent = null;
		root.parentArrtibuteValue = null;
		// root.attributeValues = getAttributeValues( maxGainIndex );
		root.attributeName = this.dataSet.getAttribute( maxGainIndex );
		// root.children = new ID3Node[root.attributeValues.length];
		insertTree( this.data, root );
		return root;
	}

	/**
	 * (gain, index)
	 * 
	 * @param array
	 * @return
	 */
	public Pair<Double, Integer> getMaxGain( Object[][] array )
	{
		double maxGain = 0;
		int maxGainIndex = 0;
		for( int i = 0; i < visable.length; i++ )
		{
			if( !visable[i] )
			{
				double value = gain( array, i );
				if( maxGain < value )
				{
					maxGain = value;
					maxGainIndex = i;
				}
			}
		}
		visable[maxGainIndex] = true;
		return Pair.wrap( maxGain, maxGainIndex );
	}

	public void insertTree( Object[][] array, ID3Node parentNode )
	{
		// Object[] parentAttrValues = parentNode.attributeValues;
		Object[] parentAttrValues = this.getAttributeValues( array, parentNode.attributeName );
		parentNode.children = new ID3Node[parentAttrValues.length];
		for( int i = 0; i < parentAttrValues.length; i++ )
		{
			Object[][] pickArray = filterArray( array, parentAttrValues[i], this.dataSet
				.getAttributeIndex( parentNode.attributeName ) );
			Pair<Double, Integer> info = getMaxGain( pickArray );
			double gain = info.first();
			if( gain != 0 )
			{
				int index = info.second();
				System.out.println( "gain = " + gain + ", attribute = "
						+ this.dataSet.getAttribute( index ) );
				ID3Node currentNode = new ID3Node();
				currentNode.parent = parentNode;
				currentNode.parentArrtibuteValue = parentAttrValues[i];
				// currentNode.attributeValues = getAttributeValues( index );
				currentNode.attributeName = this.dataSet.getAttribute( index );
				// currentNode.children = new
				// ID3Node[currentNode.attributeValues.length];
				parentNode.children[i] = currentNode;
				insertTree( pickArray, currentNode );
			}
			else
			{
				ID3Node leafNode = new ID3Node();
				leafNode.parent = parentNode;
				leafNode.parentArrtibuteValue = parentAttrValues[i];
				// leafNode.attributeValues = new String[0];
				leafNode.leafValue = getLeafNodeValue( pickArray );
				leafNode.children = new ID3Node[0];
				parentNode.children[i] = leafNode;
			}
		}
	}

	/**
	 * ����ݼ�����ѡ��ĳЩ�У���Щ����ָ�����о���ָ��ֵ
	 * 
	 * @param array
	 *        ��ݼ���
	 * @param targetValue
	 *        ָ����ֵ
	 * @param targetAttrIndex
	 *        ָ������
	 * @return
	 */
	public Object[][] filterArray( Object[][] array, Object targetValue, int targetAttrIndex )
	{
		List<Object[]> list = new ArrayList<Object[]>();
		for( int i = 0; i < array.length; i++ )
		{
			Object[] row = array[i];
			if( Utils.equals2( row[targetAttrIndex], targetValue ) )
			{
				list.add( row );
			}
		}
		return list.toArray( new Object[0][] );
	}

	/**
	 * Entropy(S)
	 * 
	 * @param array
	 * @return double
	 */
	public double gain( Object[][] array, int index )
	{
		Object[] classAttrValues = getAttributeValues( array, this.classAttrIndex );
		int[] counts = new int[classAttrValues.length];
		for( int i = 0; i < counts.length; i++ )
		{
			counts[i] = 0;
		}
		for( int i = 0; i < array.length; i++ )
		{
			Object[] strs = array[i];
			for( int j = 0; j < classAttrValues.length; j++ )
			{
				if( Utils.equals2( strs[this.classAttrIndex], classAttrValues[j] ) )
				{
					counts[j]++;
				}
			}
		}
		/**
		 * Entropy(S) = S -p(I) log2 p(I)
		 */
		double entropyS = 0;
		for( int i = 0; i < counts.length; i++ )
		{
			entropyS += ID3Utils.sigma( counts[i], array.length );
		}
		Object[] arrtValues = getAttributeValues( array, index );
		/**
		 * total ((|Sv| / |S|) * Entropy(Sv))
		 */
		double sv_total = 0;
		for( int i = 0; i < arrtValues.length; i++ )
		{
			sv_total += entropySv( array, index, arrtValues[i], array.length );
		}
		return entropyS - sv_total;
	}

	/**
	 * ((|Sv| / |S|) * Entropy(Sv))
	 * 
	 * @param array
	 * @param index
	 * @param arrtibuteValue
	 * @param allTotal
	 * @return
	 */
	public double entropySv( Object[][] array, int index, Object arrtibuteValue, int allTotal )
	{
		Object[] classAttrValues = getAttributeValues( array, this.classAttrIndex );
		int[] counts = new int[classAttrValues.length];
		for( int i = 0; i < counts.length; i++ )
		{
			counts[i] = 0;
		}
		for( int i = 0; i < array.length; i++ )
		{
			Object[] strs = array[i];
			if( Utils.equals2( strs[index], arrtibuteValue ) )
			{
				for( int k = 0; k < classAttrValues.length; k++ )
				{
					if( Utils.equals2( strs[this.classAttrIndex], classAttrValues[k] ) )
					{
						counts[k]++;
					}
				}
			}
		}
		int total = 0;
		double entropySv = 0;
		for( int i = 0; i < counts.length; i++ )
		{
			total += counts[i];
		}
		for( int i = 0; i < counts.length; i++ )
		{
			entropySv += ID3Utils.sigma( counts[i], total );
		}
		return ID3Utils.getPi( total, allTotal ) * entropySv;
	}

	public Object[] getAttributeValues( Object[][] array, String attr )
	{
		return this.getAttributeValues( array, this.dataSet.getAttributeIndex( attr ) );
	}

	public Object[] getAttributeValues( Object[][] array, int column )
	{
		Set<Object> set = new HashSet<Object>();
		for( int i = 0; i < array.length; i++ )
		{
			set.add( array[i][column] );
		}
		Object[] result = new Object[set.size()];
		return set.toArray( result );
	}

	public Object getLeafNodeValue( Object[][] array )
	{
		if( array != null && array.length > 0 )
			return array[0][classAttrIndex];
		else
		{
			return ERROR;
		}
	}
}
